# 导入基础工具库
import os                   # 用于文件路径与系统操作
import time                 # 用于计时（训练耗时、数据加载耗时等）
import random               # 用于设置随机种子，保证实验可复现
import json                 # 用于读取配置文件（如模型参数、训练设置）
from tqdm import tqdm       # 用于显示进度条（当前代码未直接使用，预留扩展）

# 导入PyTorch核心库
import torch
import torch.nn as nn                   # 用于构建神经网络层
from torch.nn.utils import clip_grad_norm_  # 用于梯度裁剪，防止梯度爆炸
from torch.utils.data import DataLoader  # 用于构建数据加载器，批量处理数据
from tensorboardX import SummaryWriter   # 用于记录训练日志，可视化训练过程
from torch.optim.lr_scheduler import StepLR, MultiStepLR  # 用于学习率调度

# 导入数值计算与自定义模块
import numpy as np
from configs.opts import parser         # 从配置文件导入参数解析器（定义训练超参）
from model.main_model import weak_main_model as main_model  # 导入弱监督主模型
from utils import AverageMeter, Prepare_logger, get_and_save_args  # 导入工具类（统计、日志、参数处理）
from utils.Recorder import Recorder     # 导入记录器（备份代码、保存参数）
from dataset.AVE_dataset_weak import AVEDataset  # 导入AVE数据集类（弱监督版本）

# =================================  随机种子配置 ============================
# 设置全局随机种子，确保模型初始化、数据打乱等过程可复现
SEED = 666
random.seed(SEED)               # Python原生随机种子
np.random.seed(seed=SEED)       # NumPy随机种子
torch.manual_seed(seed=SEED)    # PyTorch CPU随机种子
torch.cuda.manual_seed(seed=SEED)  # PyTorch GPU随机种子（单卡）
torch.backends.cudnn.deterministic = True  # 强制CuDNN使用确定性算法（保证结果一致）
torch.backends.cudnn.benchmark = False     # 禁用CuDNN自动优化（避免因算法选择导致结果差异）
# =============================================================================

# 读取弱监督训练的配置文件（如模型结构参数、数据集路径等）
config_path = 'configs/weak.json'
with open(config_path) as fp:
    config = json.load(fp)
print(config)  # 打印配置信息，方便检查是否加载正确


def main():
    # 定义全局变量（跨函数使用，避免频繁传参）
    global args, logger, writer, dataset_configs
    # 定义模型性能跟踪变量（记录最佳准确率及对应epoch）
    global best_accuracy, best_accuracy_epoch
    best_accuracy, best_accuracy_epoch = 0, 0  # 初始化为0

    # 1. 加载并解析训练参数
    dataset_configs = get_and_save_args(parser)  # 从解析器获取数据集相关配置
    parser.set_defaults(**dataset_configs)       # 将数据集配置设为解析器默认值
    args = parser.parse_args()                   # 解析命令行传入的超参（如batch_size、lr等）

    # 2. 设置GPU设备（指定使用哪些GPU）
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"  # 按PCI总线ID排序GPU（避免设备序号混乱）
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu   # 从参数指定GPU（如"0"表示使用第1块GPU）

    # 3. 创建实验目录（用于保存模型、日志、代码备份）
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)  # 若目录不存在则创建

    # 若从 checkpoint 恢复训练，将实验目录设为checkpoint所在目录
    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    # 4. 初始化日志记录器（记录训练/测试过程中的关键信息）
    logger = Prepare_logger(args, eval=args.evaluate)  # eval=True表示仅测试，False表示训练

    # 打印实验基本信息（训练模式 vs 测试模式）
    if not args.evaluate:
        logger.info(f'\nCreating folder: {args.snapshot_pref}')
        logger.info('\nRuntime args\n\n{}\n'.format(json.dumps(vars(args), indent=4)))  # 打印所有超参
    else:
        logger.info(f'\nLog file will be save in {args.snapshot_pref}/Eval.log.')  # 仅测试时指定日志路径

    # ----------------------------- 数据加载模块 -----------------------------
    '''构建训练集与测试集的数据加载器'''
    # 训练集加载器：shuffle=True（打乱数据），num_workers=8（多线程加载）
    train_dataloader = DataLoader(
        AVEDataset('./data/', split='train'),  # 加载训练集（数据根路径./data/，分裂为train）
        batch_size=args.batch_size,            # 批次大小（从超参获取）
        shuffle=True,                          # 训练时打乱数据，保证随机性
        num_workers=8,                         # 8个线程并行加载数据，加速训练
        pin_memory=True                        # 锁定内存（减少CPU到GPU的数据传输耗时）
    )

    # 测试集加载器：shuffle=False（不打乱，保证结果可复现）
    test_dataloader = DataLoader(
        AVEDataset('./data/', split='test'),   # 加载测试集（分裂为test）
        batch_size=args.test_batch_size,       # 测试时批次大小（可与训练不同）
        shuffle=False,                         # 测试时不打乱数据
        num_workers=8,
        pin_memory=True
    )

    # ----------------------------- 模型配置模块 -----------------------------
    '''初始化模型、优化器、学习率调度器、损失函数'''
    # 1. 初始化主模型（传入配置文件中的模型参数）
    mainModel = main_model(config["model"])
    # 2. 多GPU并行训练（将模型包装为DataParallel，自动分配到指定GPU）
    mainModel = nn.DataParallel(mainModel).cuda()
    # 3. 获取模型可训练参数（用于优化器）
    learned_parameters = mainModel.parameters()
    # 4. 初始化Adam优化器（学习率从超参获取）
    optimizer = torch.optim.Adam(learned_parameters, lr=args.lr)
    # 5. 初始化学习率调度器（MultiStepLR：在指定epoch降低学习率）
    # 里程碑milestones=[10,20,40]：在第10、20、40个epoch后，学习率乘以gamma=0.5
    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 40], gamma=0.5)
    # 6. 初始化损失函数（均在GPU上计算）
    criterion = nn.BCEWithLogitsLoss().cuda()          # BCEWithLogitsLoss（带sigmoid的二分类损失，预留）
    criterion_event = nn.MultiLabelSoftMarginLoss().cuda()  # 多标签软边际损失（用于事件分类）

    # ----------------------------- 恢复训练模块 -----------------------------
    '''从checkpoint恢复训练（若指定resume路径）'''
    if os.path.isfile(args.resume):
        logger.info(f"\nLoading Checkpoint: {args.resume}\n")  # 打印加载的checkpoint路径
        # 加载模型参数（注意：DataParallel的模型参数需对应）
        mainModel.load_state_dict(torch.load(args.resume))
    # 若指定了resume路径但文件不存在，抛出错误
    elif args.resume != "" and (not os.path.isfile(args.resume)):
        raise FileNotFoundError(f"Checkpoint file not found: {args.resume}")

    # ----------------------------- 仅测试模式 -----------------------------
    '''若为仅测试模式（args.evaluate=True），直接运行测试并返回'''
    if args.evaluate:
        logger.info(f"\nStart Evaluation..")
        # 调用测试函数（eval_only=True表示仅测试，不记录训练日志）
        validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch=0, eval_only=True)
        return  # 测试结束后退出程序

    # ----------------------------- 日志与代码备份 -----------------------------
    '''初始化TensorBoard日志与代码备份工具'''
    writer = SummaryWriter(args.snapshot_pref)  # TensorBoard日志写入器（保存路径为实验目录）
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")  # 代码备份器（忽略Exps/目录）
    recorder.writeopt(args)  # 将超参保存到实验目录，便于后续复现

    # ----------------------------- 核心训练循环 -----------------------------
    '''开始训练与测试循环（共args.n_epoch个epoch）'''
    for epoch in range(args.n_epoch):
        # 1. 训练一个epoch：返回当前epoch的平均训练损失
        loss = train_epoch(mainModel, train_dataloader, criterion, criterion_event, optimizer, epoch)

        # 2. 测试时机：每args.eval_freq个epoch测试一次，或最后一个epoch必测
        if ((epoch + 1) % args.eval_freq == 0) or (epoch == args.n_epoch - 1):
            # 测试一个epoch：返回当前epoch的测试准确率
            acc = validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch)
            
            # 3. 保存最佳模型：若当前准确率高于历史最佳，更新最佳记录并保存模型
            if acc > best_accuracy:
                best_accuracy = acc                # 更新最佳准确率
                best_accuracy_epoch = epoch        # 更新最佳准确率对应的epoch
                # 调用保存模型函数
                save_checkpoint(
                    mainModel.state_dict(),        # 模型参数
                    top1=best_accuracy,            # 最佳准确率
                    task='WeaklySupervised',       # 任务名称（弱监督）
                    epoch=epoch + 1,               # 当前epoch（+1是因为epoch从0开始）
                )
            # 打印当前最佳准确率与对应epoch
            print("-----------------------------")
            print("best acc and epoch:", best_accuracy, best_accuracy_epoch)
            print("-----------------------------")
        
        # 4. 更新学习率：每个epoch结束后，调度器按规则调整学习率
        scheduler.step()


def train_epoch(model, train_dataloader, criterion, criterion_event, optimizer, epoch):
    '''训练单个epoch的函数：输入模型、数据加载器等，返回当前epoch的平均损失'''
    # 初始化统计指标（AverageMeter：计算平均值，用于跟踪训练过程）
    batch_time = AverageMeter()  # 跟踪每个批次的耗时
    data_time = AverageMeter()   # 跟踪数据加载的耗时
    losses = AverageMeter()      # 跟踪训练损失
    train_acc = AverageMeter()   # 跟踪训练准确率（当前代码中acc固定为0，预留扩展）
    end_time = time.time()       # 记录当前时间（用于计算耗时）

    # 1. 设置模型为训练模式（启用Dropout、BatchNorm等训练特有的层）
    model.train()
    # 2. 设置模型参数为double类型（64位浮点数，提高计算精度）
    model.double()
    # 3. 清空优化器梯度（避免上一轮梯度累积）
    optimizer.zero_grad()

    # 遍历训练集的每个批次
    for n_iter, batch_data in enumerate(train_dataloader):
        # 更新数据加载耗时（当前时间 - 上一批次结束时间）
        data_time.update(time.time() - end_time)

        # ----------------------------- 数据预处理 -----------------------------
        '''从批次数据中提取视觉特征、音频特征、标签'''
        visual_feature, audio_feature, labels = batch_data  # 解构批次数据
        # 将标签转换为double类型并转移到GPU（与模型参数类型一致）
        labels = labels.double().cuda()

        # ----------------------------- 模型前向传播 -----------------------------
        '''将特征输入模型，获取输出'''
        # 模型输出：is_event_scores（是否为事件的得分）、raw_logits（原始logits）、event_scores（事件分类得分）
        is_event_scores, raw_logits, event_scores = model(visual_feature, audio_feature)

        # ----------------------------- 损失计算 -----------------------------
        '''计算训练损失（当前仅使用事件分类损失）'''
        # 多标签事件分类损失：event_scores（模型输出）与labels（真实标签）的损失
        loss_event_class = criterion_event(event_scores, labels)
        loss = loss_event_class  # 总损失（当前仅包含事件分类损失，预留扩展其他损失）

        # ----------------------------- 梯度反向传播 -----------------------------
        '''计算梯度并更新模型参数'''
        loss.backward()  # 反向传播：计算所有可训练参数的梯度

        # ----------------------------- 准确率计算（预留） -----------------------------
        '''当前代码中训练准确率固定为0，需根据实际任务补充计算逻辑'''
        acc = torch.tensor([0])  # 占位符：实际使用时需替换为真实准确率计算
        # 更新训练准确率统计（acc.item()为当前批次准确率，样本数为批次大小×10）
        train_acc.update(acc.item(), visual_feature.size(0) * 10)

        # ----------------------------- 梯度裁剪 -----------------------------
        '''防止梯度爆炸：若指定了梯度裁剪阈值，裁剪所有参数的梯度范数'''
        if args.clip_gradient is not None:
            total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)
            # 可选：打印梯度裁剪信息（当前注释，预留扩展）
            # if total_norm > args.clip_gradient:
            #     logger.info(f'Clipping gradient: {total_norm} with coef {args.clip_gradient/total_norm}.')

        # ----------------------------- 参数更新 -----------------------------
        '''根据梯度更新模型参数，并清空梯度'''
        optimizer.step()  # 优化器更新参数
        optimizer.zero_grad()  # 清空梯度，准备下一批次

        # ----------------------------- 统计指标更新 -----------------------------
        '''更新损失、批次耗时等统计指标'''
        # 更新损失：loss.item()为当前批次损失，样本数为批次大小×10
        losses.update(loss.item(), visual_feature.size(0) * 10)
        # 更新批次耗时（当前时间 - 上一批次结束时间）
        batch_time.update(time.time() - end_time)
        end_time = time.time()  # 更新当前时间，用于下一批次计算

        # ----------------------------- TensorBoard日志记录 -----------------------------
        '''记录每个迭代步骤的损失（用于可视化）'''
        # 日志索引：epoch×训练集批次总数 + 当前迭代次数 + 1（避免索引从0开始）
        writer.add_scalar('Train_data/loss', losses.val, epoch * len(train_dataloader) + n_iter + 1)

        # ----------------------------- 终端日志打印（预留） -----------------------------
        '''每args.print_freq个迭代步骤，打印一次训练信息（当前注释，预留扩展）'''
        # if n_iter % args.print_freq == 0:
        #     logger.info(
        #         f'Train Epoch: [{epoch}][{n_iter}/{len(train_dataloader)}]\t'
        #         f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
        #         f'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
        #         f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
        #         f'Prec@1 {train_acc.val:.3f} ({train_acc.avg: .3f})'
        #     )

    # ----------------------------- 记录epoch级损失 -----------------------------
    '''记录当前epoch的平均损失（用于TensorBoard可视化）'''
    writer.add_scalar('Train_epoch_data/epoch_loss', losses.avg, epoch)

    # 返回当前epoch的平均训练损失
    return losses.avg


@torch.no_grad()  # 装饰器：禁用梯度计算（测试阶段无需反向传播，节省内存）
def validate_epoch(model, test_dataloader, criterion, criterion_event, epoch, eval_only=False):
    """
    测试/验证单个epoch的函数：计算模型在测试集上的准确率，可选记录日志
    参数说明：
        model: 待测试的训练模型
        test_dataloader: 测试集的数据加载器（批量输出数据）
        criterion: 二分类损失函数（当前测试阶段未使用，预留接口）
        criterion_event: 多标签分类损失函数（当前测试阶段未使用，预留接口）
        epoch: 当前训练轮次（用于日志记录，区分不同epoch的测试结果）
        eval_only: 是否为“仅测试模式”（True=仅输出测试结果，False=需记录TensorBoard日志）
    返回值：
        accuracy.avg: 整个测试集的平均准确率（百分比形式）
    """
    # 初始化统计指标计算器（AverageMeter为自定义工具类，用于计算平均值）
    batch_time = AverageMeter()  # 统计每个批次的处理耗时（从数据加载到结果计算）
    data_time = AverageMeter()   # 统计每个批次的数据加载耗时
    losses = AverageMeter()      # 统计测试损失（当前未使用，预留接口）
    accuracy = AverageMeter()    # 统计测试准确率（核心指标）
    end_time = time.time()       # 记录当前时间（用于计算后续耗时）

    # 测试阶段模型必备设置
    model.eval()  # 将模型切换为“评估模式”
    # 作用：禁用训练阶段特有的层（如Dropout随机失活、BatchNorm动态更新均值/方差），保证测试结果稳定
    model.double()  # 将模型参数类型设为double（64位浮点数），与输入特征类型匹配

    # 遍历测试集的所有批次（逐批次计算准确率）
    for n_iter, batch_data in enumerate(test_dataloader):
        # 更新“数据加载耗时”：当前时间 - 上一批次结束时间
        data_time.update(time.time() - end_time)

        # ----------------------------- 1. 提取并预处理测试数据 -----------------------------
        visual_feature, audio_feature, labels = batch_data  # 从加载器获取单批次数据
        # 拆解数据：
        # visual_feature: 视觉特征（如视频帧提取的特征）
        # audio_feature: 音频特征（如音频片段提取的特征）
        # labels: 真实标签（对应视听事件的类别标注）
        labels = labels.cuda()  # 将标签转移到GPU（与模型所在设备一致）
        bs = visual_feature.size(0)  # 获取当前批次的样本数量（用于后续准确率统计）

        # ----------------------------- 2. 模型前向传播（核心测试步骤） -----------------------------
        # 输入特征到模型，输出预测结果（无梯度计算，因@torch.no_grad()装饰器）
        is_event_scores, raw_logits, event_scores = model(visual_feature, audio_feature)
        # 拆解模型输出：
        # is_event_scores: “是否存在事件”的预测得分（用于判断片段是否为有效事件）
        # raw_logits: 原始预测对数概率（未经过激活函数）
        # event_scores: 事件类别的预测得分（用于分类具体事件类型）

        # ----------------------------- 3. 计算测试准确率 -----------------------------
        # 调用自定义准确率计算函数，传入模型输出和真实标签
        # 核心逻辑：对比预测类别与真实类别，统计正确预测的比例
        acc = compute_accuracy_supervised(is_event_scores, raw_logits, labels)
        # 更新全局准确率统计：acc.item()为当前批次准确率，bs为当前批次样本数（权重）
        accuracy.update(acc.item(), bs)

        # ----------------------------- 4. 更新批次耗时统计 -----------------------------
        batch_time.update(time.time() - end_time)  # 更新“批次总耗时”
        end_time = time.time()  # 重置时间戳，为下一批次计时做准备

        # ----------------------------- 5. 打印批次测试日志（预留接口） -----------------------------
        # 每间隔args.print_freq个批次，打印当前批次的测试进度（当前注释，需结合全局参数启用）
        # if n_iter % args.print_freq == 0:
        #     logger.info(
        #         f'Test Epoch [{epoch}][{n_iter}/{len(test_dataloader)}]\t'
        #         f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'  # 当前批次耗时/平均耗时
        #         f'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'    # 当前加载耗时/平均加载耗时
        #         f'Prec@1 {accuracy.val:.3f} ({accuracy.avg:.3f})'     # 当前批次准确率/平均准确率
        #     )

    # ----------------------------- 6. 记录全局测试日志 -----------------------------
    # 非“仅测试模式”时，将平均准确率写入TensorBoard（用于可视化训练趋势）
    if not eval_only:
        writer.add_scalar('Val_epoch/Accuracy', accuracy.avg, epoch)

    # 打印整个测试集的最终准确率（通过logger写入日志文件，同时在终端显示）
    logger.info(
        f'**************************************************************************\t'
        f"\tEvaluation results (acc): {accuracy.avg:.4f}%."  # 保留4位小数的准确率
    )

    # 返回整个测试集的平均准确率（用于训练主循环判断是否保存“最佳模型”）
    return accuracy.avg


def compute_accuracy_supervised(is_event_scores, event_scores, labels):
    """
    监督式准确率计算函数（适配弱监督视听事件检测任务）
    核心逻辑：先判断“是否为事件”，再分类事件类型，最后对比真实标签统计正确数
    参数：
        is_event_scores: “是否存在事件”的预测得分
        event_scores: 事件类别的预测得分
        labels: 真实标签（维度与预测结果匹配）
    返回：
        acc: 当前批次的准确率（百分比形式）
    """
    # 步骤1：从真实标签中提取“事件类别索引”（labels为多维度，取最后一维的最大值索引）
    _, targets = labels.max(-1)  # targets: 真实事件类别索引（形状与预测结果对齐）

    # 步骤2：判断“是否为事件”（二分类）
    is_event_scores = is_event_scores.sigmoid()  # sigmoid激活：将得分转为0-1概率
    scores_pos_ind = is_event_scores > 0.5  # 概率>0.5视为“存在事件”（正样本）
    scores_mask = scores_pos_ind == 0       # 概率≤0.5视为“无事件”（负样本，生成掩码）

    # 步骤3：对“存在事件”的样本，预测具体事件类别
    _, event_class = event_scores.max(-1)  # 取事件得分的最大值索引（预测类别）
    pred = scores_pos_ind.long()  # 将“是否为事件”的布尔值转为长整型（1=有事件，0=无事件）
    pred *= event_class[:, None]  # 有事件的位置：填充预测类别；无事件的位置：保持0

    # 步骤4：对“无事件”的样本，标记为背景类别（28为自定义背景类别索引）
    pred[scores_mask] = 28  # 无事件的位置，预测类别设为28（与真实标签的背景类别对应）

    # 步骤5：计算准确率
    correct = pred.eq(targets)  # 逐元素对比预测值与真实值（True=正确，False=错误）
    correct_num = correct.sum().double()  # 统计正确预测的总数量（转为double避免整数溢出）
    acc = correct_num * (100. / correct.numel())  # 准确率=正确数/总样本数 ×100（百分比）

    return acc


def compute_accuracy_weak(event_scores, labels):
    """
    弱监督准确率计算函数（备用接口，当前主程序未使用）
    逻辑：直接对事件得分取最大值索引，与真实标签对比（不区分“是否为事件”）
    """
    _, pred = event_scores.max(-1)  # 预测类别：取事件得分最大值索引
    pred = pred.transpose(1, 0).contiguous()  # 转置维度，确保与标签形状匹配
    _, target = labels.max(-1)  # 真实类别：取标签最大值索引
    correct = pred.eq(target)  # 对比预测与真实值
    correct_num = correct.sum().double()  # 统计正确数
    acc = correct_num * (100. / correct.numel())  # 计算准确率（百分比）
    return acc


def save_checkpoint(state_dict, top1, task, epoch):
    """
    保存模型 checkpoint（用于保存“最佳模型”或定期备份）
    参数：
        state_dict: 模型的参数字典（mainModel.state_dict()）
        top1: 模型对应的准确率（用于文件名标识）
        task: 任务名称（此处为“WeaklySupervised”，弱监督任务）
        epoch: 当前训练轮次（用于文件名标识）
    """
    # 构建模型保存路径与文件名（包含关键信息，便于后续识别）
    # 格式：实验目录/模型_轮次_准确率_任务_最佳模型.pth.tar
    model_name = f'{args.snapshot_pref}/model_epoch_{epoch}_top1_{top1:.3f}_task_{task}_best_model.pth.tar'
    torch.save(state_dict, model_name)  # 保存参数字典到指定路径
    # 注：仅保存参数（state_dict），不保存完整模型，节省存储空间且便于跨环境加载


# ----------------------------- 程序执行入口 -----------------------------
if __name__ == '__main__':
    # 当脚本被直接运行时，调用main()函数启动整个训练/测试流程
    # 作用：若脚本被作为模块导入（如import），则不执行main()，避免冲突
    main()